#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) JoinQuant Development Team
# Author: Huayong Kuang <kuanghuayong@joinquant.com>

import socket

import socks

from thrift_gen.hello import HelloWorld  # 引入客户端类

from thrift import Thrift
from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol


class TProxySocket(TSocket.TSocket):
    """代理连接

    参考：https://github.com/totally/thrift_goodies/blob/master/transport.py
    """

    def __init__(self, proxy_host=None, proxy_port=None, *args, **kwargs):
        self.proxy_host = proxy_host
        self.proxy_port = proxy_port or 1080
        TSocket.TSocket.__init__(self, *args, **kwargs)

    def _do_open(self, family, socktype):
        if not self.proxy_host:
            return socket.socket(family, socktype)

        sock = socks.socksocket(family, socktype)
        sock.set_proxy(socks.SOCKS5, self.proxy_host, self.proxy_port)
        return sock


try:
    # 建立 socket
    # transport = TSocket.TSocket('localhost', 9090)
    transport = TProxySocket('localhost', 8080, 'offline2', 9090)
    # 选择传输层，这块要和服务端的设置一致
    transport = TTransport.TBufferedTransport(transport)
    # 选择传输协议，这个也要和服务端保持一致，否则无法通信
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    # 创建客户端
    client = HelloWorld.Client(protocol)
    transport.open()

    print("client - ping")
    print("server - " + client.ping())

    print("client - say")
    msg = client.say("Hello!")
    print("server - " + msg)
    # 关闭传输
    transport.close()
    # 捕获异常
except Thrift.TException as ex:
    print("%s" % (ex.message))
